# -*- coding: utf-8 -*-  
'''
crt-tplinker模型
    

@author: luoyi
Created on 2021年7月2日
'''
import tensorflow as tf

import utils.conf as conf
import utils.poses as pos


#    定义状态转移矩阵
class TransferLayer(tf.keras.layers.Layer):
    def __init__(self,
                 name='transfer_layer',
                 pos_size=pos.bio_tag_size(),
                 **kwargs):
        super(TransferLayer, self).__init__(name, **kwargs)
        
        self._pos_size = pos_size
        pass
    
    def build(self, input_shape):
        #    定义状态转移矩阵
        self._transfer = self.add_weight(name='crf_transfer', 
                                         shape=(self._pos_size, self._pos_size), 
                                         dtype=tf.float32, 
                                         initializer=tf.initializers.he_normal(), 
                                         trainable=True)
        
        super(TransferLayer, self).build(input_shape)
        pass
    
    def call(self, inputs, **kwargs):
        #    transfer[i, j]表示第i个词性转到第j个词性的概率，第i行概率和为1
        self._transfer = tf.nn.softmax(self._transfer)
        
        return inputs
    
    #    取状态转移矩阵
    def get_transfer(self):
        return self._transfer
    pass


#    定义fc层
class FCLayer(tf.keras.layers.Layer):
    def __init(self,
               name='fc_layer',
               pos_size=len(pos.bio_tag_size()),
               **kwargs):
        super(FCLayer, self).__init__(name, **kwargs)
        
        self._pos_size - pos_size
        pass
    
    def build(self, input_shape):
        #    定义fc层
        self._fc = tf.keras.layers.Dense(units=self._pos_size, 
                                         kernel_initializer=tf.initializers.he_normal(), 
                                         bias_initializer=tf.initializers.he_normal(), 
                                         )
        #    fc层输出
        self._out = None
        
        super(FCLayer, self).build(input_shape)
        pass
    
    def call(self, inputs, **kwargs):
        #    过fc层
        out = self._fc(inputs)
        self._out = out
        
        return out
    
    #    取crf_fc_layer层输出
    def get_out(self):
        '''
            @return: Tensor(batch_size, max_sen_len, pos_size)
        '''
        return self._out
    
    pass